#include "kernel_operator.h"
using namespace AscendC;

class KernelAdd {
public:
    __aicore__ inline KernelAdd() {}
    
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength) {
        this->x = x;
        this->y = y;
        this->z = z;
        this->totalLength = totalLength;
        this->blockLength = totalLength / GetBlockNum();
        this->tileNum = this->blockLength / BUFFER_SIZE;
        
        xBuf.SetBuffer(BUFFER_SIZE);
        yBuf.SetBuffer(BUFFER_SIZE);
        zBuf.SetBuffer(BUFFER_SIZE);
    }
    
    __aicore__ inline void Process() {
        for(uint32_t i = 0; i < this->tileNum; i++) {
            // 数据搬运
            DataCopy(xBuf, this->x + i * BUFFER_SIZE, BUFFER_SIZE);
            DataCopy(yBuf, this->y + i * BUFFER_SIZE, BUFFER_SIZE);
            
            // 向量加法
            Add(zBuf, xBuf, yBuf, BUFFER_SIZE);
            
            // 结果写回
            DataCopy(this->z + i * BUFFER_SIZE, zBuf, BUFFER_SIZE);
        }
    }
    
private:
    GM_ADDR x, y, z;
    uint32_t totalLength;
    uint32_t blockLength;
    uint32_t tileNum;
    LocalTensor<uint8_t> xBuf, yBuf, zBuf;
};

extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, uint32_t totalLength) {
    KernelAdd op;
    op.Init(x, y, z, totalLength);
    op.Process();
}